#pragma once
#include <vector>
#include <unordered_map>
#include <memory>
#include <cassert>

#include "include/PinPolygon.h"

struct SteinerNode : public std::enable_shared_from_this<SteinerNode> {
	PinPolygon* pin;  // 指向独立创建的PinPolygon对象
	Point position;
	double cost;
	std::weak_ptr<SteinerNode> parent;  // 使用weak_ptr避免循环引用
	std::vector<std::shared_ptr<SteinerNode>> children;  // 使用shared_ptr
	bool is_original_pin;  // 标记是否是原始输入点

	SteinerNode(PinPolygon* p, bool original = false, double c = 0,
		std::shared_ptr<SteinerNode> par = nullptr)
		: pin(p), position(p->m_center), cost(c), is_original_pin(original) {
		assert(p != nullptr && "PinPolygon must not be null for SteinerNode");
		if (par) {
			parent = par;
		}
	}

	// 安全的析构函数
	~SteinerNode() {}

	// 添加子节点
	void addChild(std::shared_ptr<SteinerNode> child) {
		if (child && child.get() != this) {
			// 如果子节点已经有父节点，先从原父节点移除
			if (auto parent_ptr = child->parent.lock()) {
				parent_ptr->removeChild(child);
			}
			children.emplace_back(child);
			child->parent = weak_from_this();  // 使用weak_ptr
		}
	}

	// 移除子节点
	void removeChild(std::shared_ptr<SteinerNode> child) {
		if (!child) return;

		auto it = std::find(children.begin(), children.end(), child);
		if (it != children.end()) {
			children.erase(it);
			child->parent.reset();
		}
	}

	// 断开与父节点的连接
	void disconnectFromParent() {
		if (auto parent_ptr = parent.lock()) {
			parent_ptr->removeChild(shared_from_this());
		}
	}
};

class SteinerTreeSolver {
public:
	using SteinerTreePtr = std::shared_ptr<SteinerNode>;
	using SteinerTreeMap = std::unordered_map<std::string, SteinerTreePtr>;

	SteinerTreeSolver() {};

	~SteinerTreeSolver() {
		for (auto& [net, polys] : m_extraPolygons) {
			for (PinPolygon* p : polys) {
				delete p;          // 补上缺失的释放
				p = nullptr;
			}
		}
	}

	// 计算所有net的斯坦纳树
	void computeSteinerTrees(
		std::unordered_map<std::string, std::vector<PinPolygon>>& nets,
		std::unordered_map<std::string, Via*>& net_vias,
		SteinerTreeMap& forest_roots, 
		std::vector<std::vector<Line>>& flyLines);

private:
	std::unordered_map<std::string, std::vector<PinPolygon*>> m_extraPolygons;

	SteinerTreePtr computeMST(const std::vector<PinPolygon*>& original_pins);
	void generateFlyLinesFromTree(SteinerTreePtr root, std::vector<Line>& flyLines);
};


/*


	SteinerTreePtr optimizeSteinerTree(SteinerTreePtr root, const std::string& net_name, const Via& via_template);
	void collectAllNodes(SteinerTreePtr root, std::vector<SteinerTreePtr>& nodes);
	void getMidPoint(const std::vector<SteinerTreePtr>& all_nodes, Point& center);
	PinPolygon* createSteinerPoint(const Point& position, const std::vector<SteinerTreePtr>& neighbors, const Via& via_template);

*/